import os
from PySide6.QtWidgets import (
    QWidget, QVBoxLayout, QFormLayout, QHBoxLayout,
    QPushButton, QFileDialog, QLineEdit, QComboBox,
    QMessageBox, QPlainTextEdit, QGroupBox, QCheckBox,
    QListWidget, QListWidgetItem, QLabel
)
from PySide6.QtCore import Qt

from src.config import app_config
from src.processing.video_composition_worker import VideoCompositionWorker
from src.utils.logger import get_logger
from src.utils.formatter import get_track_title

from src.ui.i18n import _


class CompositionTab(QWidget):
    def __init__(self, parent=None):
        super().__init__(parent)
        self.worker = None
        self.logger = get_logger(__name__)

        self._setup_ui()
        self._connect_signals()
        self.retranslate_ui()

    def _setup_ui(self):
        main_layout = QVBoxLayout(self)
        main_layout.setContentsMargins(10, 15, 10, 10)
        main_layout.setSpacing(15)

        # --- Top Layout for Inputs and Options ---
        top_container = QWidget()
        top_layout = QHBoxLayout(top_container)
        top_layout.setContentsMargins(0, 0, 0, 0)
        main_layout.addWidget(top_container)

        # --- File Inputs Group (Left Column) ---
        self.file_inputs_group = QGroupBox(_("File Inputs"))
        file_inputs_layout = QFormLayout(self.file_inputs_group)
        file_inputs_layout.setSpacing(10)
        # Give more space to file inputs
        top_layout.addWidget(self.file_inputs_group, 2)

        self.video_path_edit = QLineEdit()
        self.video_path_button = QPushButton()
        video_file_layout = QHBoxLayout()
        video_file_layout.addWidget(self.video_path_edit)
        video_file_layout.addWidget(self.video_path_button)
        self.video_path_label = QLabel()
        file_inputs_layout.addRow(self.video_path_label, video_file_layout)

        # Audio files list
        self.audio_list_widget = QListWidget()
        audio_buttons_layout = QVBoxLayout()
        self.add_audio_button = QPushButton()
        self.remove_audio_button = QPushButton()
        self.set_main_audio_button = QPushButton()
        audio_buttons_layout.addWidget(self.add_audio_button)
        audio_buttons_layout.addWidget(self.remove_audio_button)
        audio_buttons_layout.addWidget(self.set_main_audio_button)
        audio_buttons_layout.addStretch()

        audio_files_layout = QHBoxLayout()
        audio_files_layout.addWidget(self.audio_list_widget, 1)
        audio_files_layout.addLayout(audio_buttons_layout)
        self.audio_files_label = QLabel()
        file_inputs_layout.addRow(self.audio_files_label, audio_files_layout)

        # Subtitle files list
        self.subtitle_list_widget = QListWidget()
        subtitle_buttons_layout = QVBoxLayout()
        self.add_subtitle_button = QPushButton()
        self.remove_subtitle_button = QPushButton()
        self.set_main_subtitle_button = QPushButton()
        subtitle_buttons_layout.addWidget(self.add_subtitle_button)
        subtitle_buttons_layout.addWidget(self.remove_subtitle_button)
        subtitle_buttons_layout.addWidget(self.set_main_subtitle_button)
        subtitle_buttons_layout.addStretch()

        subtitle_files_layout = QHBoxLayout()
        subtitle_files_layout.addWidget(self.subtitle_list_widget, 1)
        subtitle_files_layout.addLayout(subtitle_buttons_layout)
        self.subtitle_files_label = QLabel()
        file_inputs_layout.addRow(
            self.subtitle_files_label, subtitle_files_layout)

        # --- Composition Options Group (Right Column) ---
        self.composition_group = QGroupBox(_("Composition Options"))
        composition_layout = QFormLayout(self.composition_group)
        composition_layout.setSpacing(10)
        top_layout.addWidget(self.composition_group, 1)

        self.keep_audio_checkbox = QCheckBox()
        self.keep_audio_checkbox.setChecked(True)

        self.remove_vocals_checkbox = QCheckBox()

        self.subtitle_mode_combo = QComboBox()
        self.subtitle_mode_combo.addItems(
            ["Soft Subtitles", "Hard Subtitles", "No Subtitles"])


        self.subtitle_mode_label = QLabel()
        self.extend_video_checkbox = QCheckBox()
        composition_layout.addRow(
            self.subtitle_mode_label, self.subtitle_mode_combo)
        composition_layout.addRow(self.keep_audio_checkbox)
        composition_layout.addRow(self.remove_vocals_checkbox)
        composition_layout.addRow(self.extend_video_checkbox)

        # --- Log and Button ---
        self.log_edit = QPlainTextEdit()
        self.log_edit.setReadOnly(True)
        main_layout.addWidget(self.log_edit)

        self.process_button = QPushButton()
        main_layout.addWidget(self.process_button)

    def _connect_signals(self):
        self.video_path_button.clicked.connect(self.select_video_file)
        self.add_audio_button.clicked.connect(self.add_audio_files)
        self.remove_audio_button.clicked.connect(
            lambda: self.remove_selected_item(self.audio_list_widget))
        self.set_main_audio_button.clicked.connect(
            lambda: self.set_main_item(self.audio_list_widget))
        self.add_subtitle_button.clicked.connect(self.add_subtitle_files)
        self.remove_subtitle_button.clicked.connect(
            lambda: self.remove_selected_item(self.subtitle_list_widget))
        self.set_main_subtitle_button.clicked.connect(
            lambda: self.set_main_item(self.subtitle_list_widget))
        self.process_button.clicked.connect(self.start_composition)

    def retranslate_ui(self):
        self.file_inputs_group.setTitle(_("File Inputs"))
        self.video_path_button.setText(_("Select Video File..."))
        self.video_path_label.setText(_("Video File:"))
        self.audio_list_widget.setToolTip(
            _("Add one or more dubbing audio files. Use 'Set as Main' to specify the default audio track."))
        self.add_audio_button.setText(_("Add Audio..."))
        self.remove_audio_button.setText(_("Remove Selected"))
        self.set_main_audio_button.setText(_("Set as Main ★"))
        self.audio_files_label.setText(_("Dubbing Files:"))
        self.subtitle_list_widget.setToolTip(
            _("Add one or more subtitle files. Use 'Set as Main' to specify the default subtitle track."))
        self.add_subtitle_button.setText(_("Add Subtitles..."))
        self.remove_subtitle_button.setText(_("Remove Selected"))
        self.set_main_subtitle_button.setText(_("Set as Main ★"))
        self.subtitle_files_label.setText(_("Subtitle Files:"))
        self.composition_group.setTitle(_("Composition Options"))
        self.keep_audio_checkbox.setText(
            _("Keep original audio track as a secondary track"))
        self.remove_vocals_checkbox.setText(
            _("Attempt to mix background sound with dubbing (Experimental)"))
        self.remove_vocals_checkbox.setToolTip(
            _("Tries to remove vocals from the original audio and mix the background sound with the dubbing into the default audio track.\n" 
              "This feature uses the demucs library for processing, which may significantly increase the processing time.\n" 
              "Effectiveness varies depending on the video's audio source."))
        self.subtitle_mode_combo.setItemText(0, _("Soft Subtitles"))
        self.subtitle_mode_combo.setItemText(1, _("Hard Subtitles"))
        self.subtitle_mode_combo.setItemText(2, _("No Subtitles"))
        self.subtitle_mode_label.setText(_("Subtitle Mode:"))
        self.extend_video_checkbox.setText(_("Extend video for long dubs"))
        self.extend_video_checkbox.setToolTip(
            _("When the dubbing duration is longer than the subtitle duration, extend the last frame of the video to match the dubbing duration.\nThis will re-encode the extended video segments, which may take longer.\nIf unchecked, the video and audio will be out of sync."))
        self.process_button.setText(_("Start Composition"))

    def select_video_file(self):
        last_dir = app_config.get(
            "paths.last_video_dir", os.path.expanduser("~"))
        file_path, _filter = QFileDialog.getOpenFileName(self, _(
            "Select Video File"), last_dir, _("Video Files (*.mp4 *.avi *.mkv *.mov)"))
        if file_path:
            self.video_path_edit.setText(file_path)
            app_config.set(
                "paths.last_video_dir", os.path.dirname(file_path), emit_signal=False)
            app_config.save(emit_signal=False)

    def add_audio_files(self):
        last_dir = app_config.get(
            "paths.last_dubbed_audio_dir", os.path.expanduser("~"))
        files, _filter = QFileDialog.getOpenFileNames(self, _(
            "Select one or more dubbing files"), last_dir, _("Audio Files (*.m4a *.mp3 *.wav)"))
        if files:
            for file_path in files:
                title = get_track_title(file_path)
                item = QListWidgetItem(title)
                item.setData(Qt.UserRole, file_path) # Store full path in data role
                self.audio_list_widget.addItem(item)
            if self.audio_list_widget.count() == 1:
                self.set_main_item(self.audio_list_widget,
                                   self.audio_list_widget.item(0))
            app_config.set(
                "paths.last_dubbed_audio_dir", os.path.dirname(files[0]), emit_signal=False)
            app_config.save(emit_signal=False)

    def add_subtitle_files(self):
        last_dir = app_config.get("paths.last_srt_dir", os.path.expanduser("~"))
        files, _filter = QFileDialog.getOpenFileNames(
            self, _("Select one or more SRT files"), last_dir, _("SRT Files (*.srt)"))
        if files:
            for file_path in files:
                title = get_track_title(file_path)
                item = QListWidgetItem(title)
                item.setData(Qt.UserRole, file_path) # Store full path in data role
                self.subtitle_list_widget.addItem(item)
            if self.subtitle_list_widget.count() == 1:
                self.set_main_item(self.subtitle_list_widget,
                                   self.subtitle_list_widget.item(0))
            app_config.set("paths.last_srt_dir", os.path.dirname(files[0]), emit_signal=False)
            app_config.save(emit_signal=False)

    def remove_selected_item(self, list_widget):
        for item in list_widget.selectedItems():
            list_widget.takeItem(list_widget.row(item))

    def set_main_item(self, list_widget, item=None):
        if not item:
            item = list_widget.currentItem()
        if not item:
            return

        # Remove star from all other items
        for i in range(list_widget.count()):
            it = list_widget.item(i)
            if it.text().endswith(" ★"):
                original_text = it.text()[:-2]
                it.setText(original_text)

        # Add star to the selected item if it doesn't have one
        if not item.text().endswith(" ★"):
            item.setText(f"{item.text()} ★")

    def _get_files_from_list(self, list_widget):
        files = []
        main_file = None
        for i in range(list_widget.count()):
            item = list_widget.item(i)
            path = item.data(Qt.UserRole) # Get path from data role
            text = item.text()
            if text.endswith(" ★"):
                main_file = path
                # Main file is placed at the beginning of the list
                files.insert(0, path)
            else:
                files.append(path)
        return files, main_file

    def start_composition(self):
        video_path = self.video_path_edit.text()
        audio_files, main_audio = self._get_files_from_list(
            self.audio_list_widget)
        subtitle_files, main_subtitle = self._get_files_from_list(
            self.subtitle_list_widget)

        if not video_path or not os.path.exists(video_path):
            QMessageBox.warning(self, _("Warning"), _(
                "Please select a valid video file."))
            return

        if not audio_files and not subtitle_files:
            QMessageBox.warning(self, _("Warning"), _(
                "Please add at least one dubbing file or one subtitle file."))
            return

        for file_path in audio_files + subtitle_files:
            if not os.path.exists(file_path):
                QMessageBox.warning(self, _("Warning"), _(
                    "File does not exist: {}").format(file_path))
                return

        self.process_button.setEnabled(False)
        self.log_edit.clear()
        self.log_edit.appendPlainText(_("Starting video composition task..."))

        keep_audio = self.keep_audio_checkbox.isChecked()
        remove_vocals = self.remove_vocals_checkbox.isChecked()
        subtitle_mode = self.subtitle_mode_combo.currentText()
        extend_video = self.extend_video_checkbox.isChecked()

        self.worker = VideoCompositionWorker(
            video_path=video_path,
            audio_files=audio_files,
            subtitle_files=subtitle_files,
            main_audio=main_audio,
            main_subtitle=main_subtitle,
            keep_original_audio=keep_audio,
            remove_vocals=remove_vocals,
            subtitle_mode=subtitle_mode,
            extend_video_for_dub=extend_video
        )
        self.worker.progress.connect(self.update_progress_ui)
        self.worker.composition_finished.connect(self.on_composition_finished)
        self.worker.error.connect(self.on_composition_error)
        self.worker.start()

    def update_progress_ui(self, message):
        self.log_edit.appendPlainText(message)

    # handle_log_message method is no longer needed as it connects directly to the main window's log_callback

    def on_composition_finished(self, message, output_path):
        self.log_edit.appendPlainText(f"\n🎉 {message}")
        self.log_edit.appendPlainText(
            _("Final video file saved at: {}").format(output_path))
        self.process_button.setEnabled(True)
        QMessageBox.information(self, _("Processing Complete"), _(
            "Video composition successful!\nOutput file: {}").format(output_path))

    def on_composition_error(self, error_message):
        self.log_edit.appendPlainText(
            f"\n❌ " + _("Error: {}").format(error_message))
        self.process_button.setEnabled(True)
        QMessageBox.critical(self, _("Processing Failed"), error_message)
